There was a great Washington Post article on social distancing. It comes with simulations to show you the effect of different measures, and I think it's fantastic #SciCom. I want to talk it over with my kids, but I want to be ready for them to ask questions like "what if we did (some cool other idea) for social distancing?" So, before I talk it over with them, I want to have some Python code to replicate it and play with. Here are my initial thoughts:
run is fast enough for like 100-1000 particles. The animation is slow. That means I want to pregenerate the trajectory, and then graph that later.Here's the python setup, using anaconda
conda create --name covid python=3 jupyter jupyterlab numpy scipy matplotlib pandas
conda activate covid
import numpy as np, scipy as sp, pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics.pairwise import euclidean_distances
from matplotlib import animation
from IPython.display import HTML
from itertools import combinations
from collections import namedtuple
#%load_ext line_profiler
"""
Units: time is measured in 12 hour chunks, but there's no strict correspondence to units here.
"""
class EffectiveArea:
"""Tells us the boundaries of the area that people can live in.
"""
def __init__(self,):
self.xmin, self.xmax = 0, 20
self.ymin, self.ymax = 0, 20
class Person:
def __init__(self,effectivearea,state='succeptible',distancing=False):
self.r = 0.15
self.days_infected = 0
self._typical_recovery_days = 14
self.distancing = distancing
if self.distancing:
self.m = 1000
else:
self.m = 1
self.state = state
self.ea = effectivearea
self.x, self.y = np.random.uniform(self.ea.xmin,self.ea.xmax), np.random.uniform(self.ea.ymin,self.ea.ymax)
if self.distancing:
self.vx,self.vy = 0,0
else:
self.vx, self.vy = np.random.normal(size=2) # Maxwell Boltzmann?
def move(self,dt,ea):
self.x, self.y = self.x + self.vx*dt, self.y + self.vy*dt
if self.state == 'infected':
if self.days_infected < self._typical_recovery_days:
self.days_infected += dt
else:
if np.random.uniform() < 0.5:
self.state = 'recovered'
self.days_infected = 0
else:
self.days_infected += dt
def collide(p1,p2):
if p1.state == 'infected' and p2.state == 'succeptible':
p2.state = 'infected'
elif p2.state == 'infected' and p1.state == 'succeptible':
p1.state = 'infected'
m1, m2 = p1.m, p2.m
r1, r2 = np.array([p1.x,p1.y]), np.array([p2.x,p2.y])
v1, v2 = np.array([p1.vx,p1.vy]), np.array([p2.vx,p2.vy])
M = m1 + m2
d = np.linalg.norm(r1 - r2)**2
u1 = v1 - 2*m2 / M * np.dot(v1-v2, r1-r2) / d * (r1 - r2)
u2 = v2 - 2*m1 / M * np.dot(v2-v1, r2-r1) / d * (r2 - r1)
p1.vx,p1.vy = u1
p2.vx,p2.vy = u2
class Universe:
def __init__(self,npeople,initial_infection_chance=0.1,
distancing=0.0):
self.npeople = npeople
self.ea = EffectiveArea()
self.dt = 0.1
self.data = None # gets set in self.run
def _state():
if np.random.uniform() < initial_infection_chance:
return 'infected'
return 'succeptible'
def _distancing():
if np.random.uniform() < distancing:
return True
return False
self.people = [Person(self.ea,_state(),_distancing()) for i in range(self.npeople)]
self.color = {'succeptible':0.5,'infected':0.0,'recovered':0.7}
self.color = {'succeptible':'lightblue','infected':'red','recovered':'green'}
def _step(self):
points = list(zip([p.x for p in self.people],[p.y for p in self.people]))
dists = euclidean_distances(points,points)
close = dists < 2*self.people[0].r
close = close.tolist()
for (i,j) in combinations(range(self.npeople),2):
if close[i][j]: # a bit faster than numpy indexing once things get big.
collide(self.people[i],self.people[j])
for p in self.people:
p.move(self.dt,self.ea)
if p.x < self.ea.xmin or p.x > self.ea.xmax:
p.vx = -p.vx
if p.y < self.ea.ymin or p.y > self.ea.ymax:
p.vy = -p.vy
def run(self,steps):
#x_coords[frame,particle_number]
x_coords = np.zeros((steps,len(self.people)))
y_coords = np.zeros((steps,len(self.people)))
state = np.zeros((steps,len(self.people)),dtype='object')
s,i,r = np.zeros(steps),np.zeros(steps),np.zeros(steps)
def pop_count(people):
s,i,r = 0,0,0
for p in people:
if p.state == 'succeptible':
s += 1
elif p.state == 'infected':
i += 1
elif p.state == 'recovered':
r += 1
return s,i,r
s[0],i[0],r[0] = pop_count(self.people)
x_coords[0] = [p.x for p in self.people]
y_coords[0] = [p.y for p in self.people]
state[0] = [p.state for p in self.people]
for step in range(1,steps):
self._step()
s[step],i[step],r[step] = pop_count(self.people)
x_coords[step] = [p.x for p in self.people]
y_coords[step] = [p.y for p in self.people]
state[step] = [p.state for p in self.people]
dtype = namedtuple('RunData',['x','y','state','s','i','r','steps'])
self.data = dtype(x_coords,y_coords,state,s,i,r,steps)
def draw(self,ax=None):
if ax is None:
fig,ax = plt.subplots(figsize=(5,5))
scat = ax.scatter([p.x for p in self.people],[p.y for p in self.people],
c = [self.color[p.state] for p in self.people],
marker='.')
return scat,
def getanim(u):
fig,ax = plt.subplots(figsize=(4,4))
left, width = 0.1, 0.85
bottom, height = 0.1, 0.65
spacing = 0.02
rect_universe = [left, bottom, width, height]
rect_trend = [left, bottom + height + spacing, width, 0.2]
ax_universe = plt.axes(rect_universe)
ax_universe.tick_params(axis='x', which='both',bottom=False,top=False,labelbottom=False)
ax_universe.axis('off')
ax_trend = plt.axes(rect_trend)
s,i,r = u.data.s,u.data.i,u.data.r
ax_trend.stackplot(range(len(s)), i, s, r, labels=['sick','healthy','recovered'],colors=['red','lightblue','green'])
scat, = u.draw(ax_universe)
def drawframe(i):
data = np.column_stack(([u.data.x[i],u.data.y[i]]))
scat.set_offsets(data)
colors = np.array([u.color[_] for _ in u.data.state[i]])
scat.set_color(colors)
_s,_i,_r = np.zeros(len(u.data.s)),np.zeros(len(u.data.s)),np.zeros(len(u.data.s))
_s[:i],_i[:i],_r[:i] = u.data.s[:i],u.data.i[:i],u.data.r[:i]
#ax_trend.stackplot(range(len(_s)), _s, _i, _r, labels=['s','i','r'],colors=['blue','green','yellow'])
#ax_trend.legend(loc='upper left')
return scat,
anim = animation.FuncAnimation(fig, drawframe, frames=u.data.steps,
interval=20, blit=False, repeat=False)
return anim
For all of these, we'll do 200 people, 200 steps, initial infection chance 10%
npeople, nsteps = 200, 200
initial_infection_chance = 0.1
%%capture
distancing = 0.0
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
HTML(anim.to_jshtml())
%%capture
distancing = 0.25
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
HTML(anim.to_jshtml())
%%capture
distancing = 0.5
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
HTML(anim.to_jshtml())
%%capture
distancing = 0.75
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
HTML(anim.to_jshtml())
%%capture
distancing = 0.90
u = Universe(npeople,distancing=distancing,initial_infection_chance=initial_infection_chance)
u.run(nsteps)
anim = getanim(u)
HTML(anim.to_jshtml())